# @Time     : 2025/4/28 11:36
# @Software : Python 3.10
# @About    :
import os.path
from pathlib import Path
from typing import IO, Union, Optional, List, Dict

from docx import Document
from docx import shared
from docx.enum.table import WD_ALIGN_VERTICAL
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.parts.document import DocumentPart
from docx.shared import Pt
from docx.styles.style import ParagraphStyle
from docx.styles.styles import Styles
from docx.text.paragraph import Paragraph

from . import tags

PARAGRAPH_ALIGNMENT_MAP = {
    "left": WD_PARAGRAPH_ALIGNMENT.LEFT,
    "right": WD_PARAGRAPH_ALIGNMENT.RIGHT,
    "center": WD_PARAGRAPH_ALIGNMENT.CENTER,
}


class DocxDocumentWriter:
    def __init__(self, docx: Optional[Union[str, IO[bytes]]] = None,
                 tag2style_map: Optional[Dict[str, str]] = None,
                 class2style_map: Optional[Dict[str, str]] = None,
                 remove_default_paragraph: bool = True):
        """
        Docx文档Writer

        Args:
            docx: **Optional[Union[str, IO[bytes]]]** *default=None* 模版文档路径
            tag2style_map:
            class2style_map:
            remove_default_paragraph: **bool** *default=True* 是否清空默认段落
        """
        if docx is None:
            docx = os.path.join(Path(__file__).parent, 'resource', 'default.docx')
        self._document = Document(docx)
        default_section = self._document.sections[0]
        self._page_view_width = shared.Cm(
            int(default_section.page_width.cm - default_section.left_margin.cm - default_section.right_margin.cm))
        self._page_view_height = shared.Cm(
            int(default_section.page_height.cm - default_section.top_margin.cm - default_section.bottom_margin.cm))
        if remove_default_paragraph:
            while len(self._document.paragraphs) > 0:
                self.remove_paragraph(self._document.paragraphs[0])

        self._document_styles = [i.name for i in self._document.styles]
        self._tag2style_map = {}
        self._class2style_map = {}
        self.update_tag2style_map(**tag2style_map or {})
        if class2style_map:
            self.update_class2style_map: (class2style_map)

    def save(self, path_or_stream: Union[str, IO[bytes]]):
        self._document.save(path_or_stream)

    @property
    def tag2style_map(self):
        return self._tag2style_map

    @property
    def class2style_map(self):
        return self._class2style_map


    def update_tag2style_map(self,
                             title: str = "Title",
                             title_1: str = "Heading 1", title_2: str = "Heading 2", title_3: str = "Heading 3",
                             title_4: str = "Heading 4", title_5: str = "Heading 5", title_6: str = "Heading 6",
                             text: str = "Text", bold_text: str = "BoldText", italic_text: str = "ItalicText",
                             paragraph: str = "Normal", quote: str = "Intense Quote",
                             unordered_list: str = "List Bullet", ordered_list: str = "List Number",
                             image: str = None, table: str = "Table Grid"):
        """
        更新样式映射

        Args:
            title:
            title_1:
            title_2:
            title_3:
            title_4:
            title_5:
            title_6:
            text:
            bold_text:
            italic_text:
            paragraph:
            quote:
            unordered_list:
            ordered_list:
            image:
            table:

        Returns:
           **DocxDocumentWriter**

        """
        self._tag2style_map.update({
            "title": title,
            "title_1": title_1,
            "title_2": title_2,
            "title_3": title_3,
            "title_4": title_4,
            "title_5": title_5,
            "title_6": title_6,
            "text": text,
            "bold_text": bold_text,
            "italic_text": italic_text,
            "paragraph": paragraph,
            "quote": quote,
            "unordered_list": unordered_list,
            "ordered_list": ordered_list,
            "image": image,
            "table": table
        })
        return self

    def update_class2style_map(self, style_map: Dict[str, str]):
        """
        更新样式映射

        Args:
            style_map: **Dict[str, str]**

        Returns:
            **DocxDocumentWriter**

        """
        for name, user_style in style_map.items():
            if user_style not in self._document_styles:
                raise ValueError(f'Can not update class to style map, style [{user_style}] not found in document')
            else:
                self._class2style_map[name] = user_style
        return self

    @property
    def document(self) -> Document:
        return self._document

    @property
    def styles(self) -> Styles:
        """A |Styles| object providing access to the styles in this document."""
        return self._document.styles

    @property
    def part(self) -> DocumentPart:
        return self._document.part

    def _check_style(self, obj: tags.Tag, style: Optional[str] = None):
        if style is not None:
            _style = style
            assert _style in self._document_styles, \
                f'Style [{_style}] not found in document styles {self._document_styles}'
        else:
            if obj.style:
                _style = self._class2style_map.get(obj.style, obj.style)
                assert _style in self._document_styles, \
                    (f"Style [{_style}] set by class name"
                     f"{'(original class name {text.style})' if _style == obj.style else ''},"
                     f" bug not found in document styles {self._document_styles}")
            else:
                _style = self._tag2style_map.get(obj.tag)
                if _style is None:
                    _style = None
                else:
                    assert _style in self._document_styles, \
                        (f'Style [{_style}] set by tag <{obj.tag}>,'
                         f'but not found in document styles {self._document_styles}')

        return _style

    def restart_numbering(self, paragraph: Paragraph, style: Union[str, ParagraphStyle] = 'List Number'):
        """
        重新编号

        Args:
            paragraph: `Paragraph`
            style: `Union[str, ParagraphStyle]`

        Returns:
            **DocxDocumentWriter**

        Reference:
            https://stackoverflow.com/questions/78683793/python-docx-restart-list-numbering

        """
        target_style_name = style.name if isinstance(style, ParagraphStyle) else style
        num_id_list_number = -1
        num_id_list_number_new = -1

        for style in self.styles:
            if style.name == target_style_name:
                if len(style.element.pPr):
                    num_id_list_number = style.element.pPr.numPr.numId.val

        # 添加新的编号链接到相同的abstractNumId，但具有startOverride并获取新的numId

        if num_id_list_number > -1:
            ct_element = self.part.numbering_part.element
            abstract_num_id = ct_element.num_having_numId(num_id_list_number).abstractNumId.val
            ct_num = ct_element.add_num(abstract_num_id)
            num_id_list_number_new = ct_num.numId
            ct_num.add_lvlOverride(0).add_startOverride(1)

        if num_id_list_number_new > -1:
            paragraph._element.pPr._add_numPr()._add_numId().val = num_id_list_number_new

        return self

    def remove_paragraph(self, paragraph: Paragraph):
        """
        删除段落

        Args:
            paragraph: **Paragraph**

        Returns:
            **DocxDocumentWriter**

        """
        paragraph._element.getparent().remove(paragraph._element)
        return self

    def add_heading(self, title: Union[tags.Heading, str], style: Optional[str] = None, return_self: bool = True):
        """
        添加标题

        Args:
            title: **Union[tags.Heading, str]** 标题内容
            style: **Optional[str]** *default=None* 样式名称
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **Paragraph**

        """
        if not hasattr(title, 'tag'):
            title = tags.Heading(title)

        _style = self._check_style(title, style)
        if _style:
            _paragraph = self._document.add_paragraph(str(title).strip(), style=_style)
        else:
            _paragraph = self._document.add_heading(str(title).strip(), level=title.level if title.level > 0 else 1)

        return self if return_self else _paragraph

    def add_paragraph(self, text: Union[tags.Texts, str], style: Optional[str] = None,
                      first_line_indent: int = 0, return_self: bool = True):
        """
        添加段落

        Args:
            text: **Union[tags.Texts, str]** 段落内容
            style: **Optional[str]** *default=None* 样式名称
            first_line_indent: **int** *default=0* 首行缩进字符
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **Paragraph**

        """
        if not hasattr(text, 'tag'):
            text = tags.Texts([tags.Text(text)])
        _style = self._check_style(text, style)
        _paragraph = self._document.add_paragraph(style=_style)
        # 如果只有一张图片，则直接插入图片，并设置居中
        if len(text) == 1 and isinstance(text[0], tags.Image):
            self.add_picture(text[0], paragraph=_paragraph)
            return self
        for part in text:
            if isinstance(part, tags.Image):
                self.add_picture(part, paragraph=_paragraph, center=False)
            elif isinstance(part, tags.BoldText):
                _paragraph.add_run(str(part)).bold = True
            elif isinstance(part, tags.ItalicText):
                _paragraph.add_run(str(part)).italic = True
            else:
                _paragraph.add_run(str(part))
        if _paragraph.style.font.size:
            _paragraph.paragraph_format.first_line_indent = Pt(_paragraph.style.font.size.pt * first_line_indent)

        return self if return_self else _paragraph

    def add_quote_paragraph(self, text: Union[tags.IntenseQuote, str], style: Optional[str] = None,
                            return_self: bool = True):
        """
        添加引用段落

        Args:
            text: **Union[tags.IntenseQuote, str]** 段落内容
            style: **Optional[str]** *default=None* 样式名称
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **Paragraph**

        """
        if not hasattr(text, 'tag'):
            text = tags.IntenseQuote([tags.Text(text)])
        _style = self._check_style(text, style)
        _quote = self._document.add_paragraph(style=_style)
        for part in text:
            if isinstance(part, tags.BoldText):
                _quote.add_run(str(part)).bold = True
            elif isinstance(part, tags.ItalicText):
                _quote.add_run(str(part)).italic = True
            else:
                _quote.add_run(str(part))
        return self if return_self else _quote

    def add_ordered_list_paragraph(self, items: Union[tags.OrderedList, List[str]], style: str = None,
                                   continue_num: bool = False, return_self: bool = True):
        """
        添加有序列表

        Args:
            items: **Union[List[tags.OrderedList], List[str]]**
            style: **Optional[str]** *default=None* 样式名称
            continue_num: **bool** *default=False* 是否继续排序序号
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **List[Paragraph]**

        """
        if not hasattr(items, 'tag'):
            items = tags.OrderedList([tags.Texts([tags.Text(item)]) for item in items])
        _style = self._check_style(items, style)
        paragraphs = []
        for _idx, item in enumerate(items):
            paragraph = self._document.add_paragraph(style=_style)
            if _idx == 0 and not continue_num:
                self.restart_numbering(paragraph, _style)
            for part in item:
                if isinstance(part, tags.BoldText):
                    paragraph.add_run(str(part)).bold = True
                elif isinstance(part, tags.ItalicText):
                    paragraph.add_run(str(part)).italic = True
                else:
                    paragraph.add_run(str(part))
            paragraphs.append(paragraph)
        return self if return_self else paragraphs

    def add_unordered_list_paragraph(self, items: Union[List[tags.UnorderedList], List[str]], style: str = None,
                                     return_self: bool = True):
        """
        添加无序列表

        Args:
            items: **Union[List[tags.UnorderedList], List[str]]**
            style: **Optional[str]** *default=None* 样式名称
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **List[Paragraph]**

        """
        if not hasattr(items, 'tag'):
            items = tags.UnorderedList([tags.Texts([tags.Text(item)]) for item in items])
        _style = self._check_style(items, style)
        paragraphs = []
        for item in items:
            paragraph = self._document.add_paragraph(style=_style)
            self.restart_numbering(paragraph, _style)
            for part in item:
                if isinstance(part, tags.BoldText):
                    paragraph.add_run(part).bold = True
                elif isinstance(part, tags.ItalicText):
                    paragraph.add_run(part).italic = True
                else:
                    paragraph.add_run(part)
            paragraphs.append(paragraph)
        return self if return_self else paragraphs

    def add_picture(self, image: tags.Image, width: Optional[int] = None, height: Optional[int] = None,
                    paragraph: Optional[Paragraph] = None, center: bool = True, return_self: bool = True):
        """
        添加图片
        
        Args:
            image: **tags.Image**
            width: **Optional[int]**  *default=None*
            height: **Optional[int]**  *default=None*
            paragraph: **Optional[Paragraph]**  *default=None* 代插入的段落
            center: **bool** *default=True* 居中对齐
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **Paragraph**

        """
        if paragraph:
            _run = paragraph.add_run("")
        else:
            paragraph = self._document.add_paragraph()
            _run = paragraph.add_run("")

        # 设置单倍行间距
        paragraph.paragraph_format.line_spacing = 1
        if center:
            paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
        _picture = _run.add_picture(image.read_image(), width, height)
        if width is None and height is None:
            if _picture.width > self._page_view_width or _picture.height > self._page_view_height:
                if _picture.width > _picture.height:
                    source_width = _picture.width
                    _picture.width = self._page_view_width
                    _picture.height = int(_picture.height * (_picture.width / source_width))
                else:
                    if int(_picture.width * (self._page_view_height / _picture.height)) < self._page_view_width:
                        source_width = _picture.width
                        _picture.width = self._page_view_width
                        _picture.height = int(_picture.height * (_picture.width / source_width))
                    else:
                        source_height = _picture.height
                        _picture.height = self._page_view_height
                        _picture.width = int(_picture.width * (_picture.height / source_height))
        return self if return_self else paragraph

    def add_table(self, table: tags.Table, style: str = None, return_self: bool = True):
        """
        添加表格
        
        Args:
            table: **tags.Table**
            style: **Optional[str]** *default=None* 样式名称
            return_self: **bool** *default=False* 返回Self还是新增对象

        Returns:
            **DocxDocumentWriter** or **Table**

        """
        _style = self._check_style(table, style)
        _table = self._document.add_table(*table.size, style=_style)
        for cell in table:
            _cell = _table.cell(cell.row, cell.col)
            _cell.text = cell.text.strip()
            # 水平位置
            _cell.paragraphs[0].alignment = PARAGRAPH_ALIGNMENT_MAP.get(cell.align, 'center')
            # 垂直位置
            _cell.vertical_alignment = WD_ALIGN_VERTICAL.CENTER
            if cell.is_header:
                _cell.paragraphs[0].runs[0].bold = True
            # 合并单元格
            if cell.colspan > 1 or cell.rowspan > 1:
                _cell.merge(_table.cell(cell.row + cell.rowspan - 1, cell.col + cell.colspan - 1))

        return self if return_self else table
